package spinal_loongarch_core132

import spinal.core._
import spinal.lib._


class TransResult extends Bundle{
    val pt   :UInt = UInt(LISA.w_ppn bits)
    val mat  :UInt = UInt(2 bits)
    val ex   :Bool = Bool()
    val ecode:UInt = LISA.Ex.Code
    val esubc:UInt = LISA.Ex.Subc
    def assignedFrom(pa:UInt,mat:UInt):this.type = {
        this.pt  := pa(pa.high downto pa.getWidth - LISA.w_ppn)
        this.mat := mat
        this.ex    := False
        this.ecode := U(0)
        this.esubc := U(0)
        this
    }
}
class TlbSrch(val isFetch:Boolean) extends Bundle with IMasterSlave{
    val va   :UInt = LISA.GPR
    val st   :Bool = if(isFetch) null else Bool()
    val valid:Bool = if(isFetch) null else Bool()
    val res:TransResult = new TransResult
    override def asMaster():Unit = {
        out(va)
        in(res)
        if(!isFetch){
            out(st,valid)
        }
    }
    def pa   :UInt = if(LISA.withMMU) res.pt @@ va(11 downto 0) else va
    def ex   :Bool = if(LISA.withMMU) res.ex    else False
    def mat  :UInt = if(LISA.withMMU) res.mat   else U"00"
    def ecode:UInt = if(LISA.withMMU) res.ecode else LISA.Ex.Code(0)
    def esubc:UInt = if(LISA.withMMU) res.esubc else LISA.Ex.Subc(0)
}
class ITlbSrch extends TlbSrch(true )
class DTlbSrch extends TlbSrch(false)
class TlbCsr extends Bundle{
    val crmd   = new CsrCRMD
    val tlbidx = new CsrTLBIDX
    val asid   = new CsrASID
    val tlbehi = new CsrTLBEHI
    val tlbelo = Vec(new CsrTLBELO, 2)
    val dmw    = Vec(new CsrDMW, 2)
    def toTlbEHi():TlbEHi = {
        val hi = new TlbEHi
        hi.e    := !tlbidx.ne
        hi.vppn := tlbehi.vppn
        hi.ps   := tlbidx.ps
        hi.asid := asid.asid
        hi.g    := tlbelo.map(_.g).reduce(_&&_)
        hi
    }
    def collect(csr:Csr){
        crmd   := csr.crmd  
        tlbidx := csr.tlbidx
        asid   := csr.asid  
        tlbehi := csr.tlbehi
        tlbelo(0) := csr.tlbelo0
        tlbelo(1) := csr.tlbelo1
        dmw(0) := csr.dmw0
        dmw(1) := csr.dmw1
    }
}
class TlbRW extends Bundle{
    val srch:Bool = Bool()
    val rd  :Bool = Bool()
    val wr  :Bool = Bool()
    val fill:Bool = Bool()
    def en:Bool = wr || fill
    def setEmpty:this.type = {
        srch := False
        rd   := False
        wr   := False
        fill := False
        this
    }
}
class TlbSrchRes extends Bundle{
    val valid = Bool()
    val hit   = Bool()
    val index = UInt(LISA.w_tlbidx bits)
    def assignedFrom(lookup:MmuLookup){
        hit   := lookup.tlb.hit
        index := Mux(lookup.tlb.hit,OHToUInt(lookup.tlb.hits),U(0))
    }
}
class TlbRdRes extends Bundle{
    val valid  = Bool()
    val e      = Bool()
    class RdData extends Bundle{
        val vppn   = UInt(19 bits)
        val asid   = UInt(LISA.w_asid bits)
        val ps     = UInt(6 bits)
        val lo     = Vec(new CsrTLBELO, 2)
        def cleared:this.type = {
            vppn := U(0)
            asid := U(0)
            ps   := U(0)
            lo.foreach(_.cleared)
            return this
        }
    }
    val dat = new RdData
    def readFrom(cam:Vec[TlbEntry],csr:TlbCsr):Unit = {
        val res:TlbEntry = cam(csr.tlbidx.index)
        val empty:RdData =(new RdData).cleared
        val rdata:RdData = new RdData
        this.e := res.hi.e
        rdata.vppn := res.hi.vppn
        rdata.asid := res.hi.asid
        rdata.ps   := res.hi.ps
        for(i <- 0 to 1){
            rdata.lo(i).assignSomeByName(res.lo(i))
            rdata.lo(i).g := res.hi.g
        }
        this.dat := Mux(res.hi.e,rdata,empty)
    }
    def wrapWrite(t:Bundle,s:Bundle,style:String=LISA.tlbrd_style):Unit = {
        style match {
        case "clear" =>
            t.assignSomeByName(s)
        case "keep" => 
            when(e){
                t.assignSomeByName(s)
            }
        }
    }
}
class TlbEHi extends Bundle{
    val vppn:UInt = UInt(19 bits)
    val ps  :UInt = UInt(6  bits)
    val g   :Bool = Bool()
    val asid:UInt = UInt(10 bits)
    val e   :Bool = Bool()
    var ps_d :Bits = null
    var ps_c :Bool = null
    var mask :Bits = null
    def decodePS():Unit = {
        if(ps_d != null)return
        ps_d = Vec(Range(0,LISA.ps_n).map((i:Int) => ps === U(LISA.psizes(i)))).asBits
        ps_c = ps_d(1)
        val w:Int = LISA.ps_max - LISA.ps_min
        mask = B("1" * (31-LISA.ps_max)) ## Mux(ps_d(0),B("1"*w),B(0))
    }
    def compare(va:UInt):Bool = {
        decodePS()
        ps_d.orR && ((va(va.high downto 13) ^ vppn).asBits & mask) === B(0)
    }
    def getSelection(va:UInt):Bool = {
        decodePS()
        Range(0,LISA.psizes.length).map((i:Int) => ps_d(i) && va(LISA.psizes(i))).reduce(_||_)
    }
}
class TlbEntry extends Bundle{
    val hi:TlbEHi      = new TlbEHi
    val lo:Vec[CsrTLBELONoG] = Vec(new CsrTLBELONoG,2)
}
class InvTlb extends Bundle{
    val en  :Bool = Bool()
    val may :Bool = Bool()
    val op  :UInt = UInt(5 bits)
    val va  :UInt = LISA.GPR
    val asid:UInt = UInt(LISA.w_asid bits)
    def getInvMask(lookup:TlbLookup):Bits = {
        val n:Int = lookup.bits_e.getWidth
        val masks :Vec[Bits] = Vec(Bits(n bits),8)
        masks(0) := B("1" * n)
        masks(1) := B("1" * n)
        masks(2) := lookup.bits_e &  lookup.bits_g
        masks(3) := lookup.bits_e & ~lookup.bits_g
        masks(4) := lookup.bits_e & ~lookup.bits_g & lookup.hits_asid
        masks(5) := lookup.bits_e & ~lookup.bits_g & lookup.hits_asid & lookup.hits_va
        masks(6) := lookup.bits_e & (lookup.bits_g | lookup.hits_asid)& lookup.hits_va
        masks(7) := B("0" * n)
        return Mux(op(4 downto 3).orR,B(0),masks(op(2 downto 0)))
    }
}
trait SomeLookup{
    val r:LookupReq
    def isFetch:Boolean = r.isFetch
    def valid :Bool = r.valid
    val found :Bool
    val result:TransResult
}
case class DaMLookup(r:LookupReq,crmd:CsrCRMD) extends Area with SomeLookup{
    val mat    = if(isFetch) crmd.datf else crmd.datm
    val found  = crmd.da
    val result = new TransResult assignedFrom(r.va,mat)
}
case class DmwLookup(r:LookupReq,csr:TlbCsr) extends Area with SomeLookup{
    val hits = csr.dmw.map(_.compare(r.va,csr.crmd.plv))
    val pa    = LISA.PA
    val mat   = UInt(2 bits)
    if(LISA.hasDmwPriority){
        when(hits(0)){
            pa  := csr.dmw(0).translate(r.va)
            mat := csr.dmw(0).mat
        }.elsewhen(hits(1)){
            pa  := csr.dmw(1).translate(r.va) 
            mat := csr.dmw(1).mat
        }.otherwise{
            pa  := U(0)
            mat := U(0)
        }
    }
    else{
        pa  := MuxOH0(hits,csr.dmw.map(_.translate(r.va))).asUInt
        mat := MuxOH0(hits,csr.dmw.map(_.mat)).asUInt
    }
    val found  = hits.asBits.orR
    val result = new TransResult assignedFrom(pa,mat)
}
case class TlbLookup(r:LookupReq,csr:TlbCsr,cam:Vec[TlbEntry]) extends Area with SomeLookup{
    val n:Int = cam.length
    def getASID:UInt = if(r.asid != null) r.asid else csr.asid.asid
    def va_comb  :UInt = r.va
    val hits_asid:Bits = cam.map(_.hi.asid === getASID).asBits
    val hits_va  :Bits = cam.map(_.hi.compare(va_comb)).asBits
    val bits_e   :Bits = cam.map(_.hi.e).asBits
    val bits_g   :Bits = cam.map(_.hi.g).asBits
    val sels_comb:Bits = cam.map(_.hi.getSelection(va_comb)).asBits
    val hits_comb:Bits = bits_e & (bits_g | hits_asid) & hits_va
    val sels     :Bits = if(valid != null) RegNextWhen(sels_comb,valid) else sels_comb
    val hits     :Bits = if(valid != null) RegNextWhen(hits_comb,valid) else hits_comb
    val is_st    :Bool = if(valid != null) RegNextWhen(r.st     ,valid) else r.st
    val va       :UInt = if(valid != null) RegNextWhen(r.va     ,valid) else r.va
    val hit      :Bool = hits.orR
    val hits_arr :Vec[Bool] = Vec(for(i <- 0 until hits.getWidth) yield hits(i))
    def select(data:IndexedSeq[Bits]):Bits = MuxOH0(hits_arr,data)
    val ps_c:Bool = select(cam.map(_.hi.ps_c.asBits)).asBool
    def getLoBits(i:Int) = Mux(sels(i),cam(i).lo(1).asBits,cam(i).lo(0).asBits)
    val lo_bvec:Vec[Bits] = Vec(for(i <- 0 until n) yield getLoBits(i))
    val lo_bits:Bits = select(lo_bvec)
    val lo:CsrTLBELONoG = new CsrTLBELONoG
    lo.assignFromBits(lo_bits)
    val is_ade :Bool = csr.crmd.plv === U(3) && va.msb
    val is_tlbr:Bool = ! hit 
    val is_inv :Bool = ! lo.v
    val is_ppi :Bool = lo.plv < csr.crmd.plv
    val is_mod :Bool = if(isFetch) null else is_st && !lo.d
    val ex_list:List[(Bool,LISA.ExFactory)] = (if(isFetch) List(
            (is_ade ,LISA.Ex.ADEF),
            (is_tlbr,LISA.Ex.TLBR),
            (is_inv ,LISA.Ex.PIF ),
            (is_ppi ,LISA.Ex.PPI )
        ) else List(
            (is_ade ,LISA.Ex.ADEM),
            (is_tlbr,LISA.Ex.TLBR),
            (is_inv ,LISA.Ex.PISL(is_st)),
            (is_ppi ,LISA.Ex.PPI ),
            (is_mod ,LISA.Ex.PME )
        )
    )
    val result = new TransResult
    val found:Bool = null
    result.pt(0, 10 bits) := Mux(ps_c,va(21 downto 12),lo.ppn(9 downto 0))
    result.pt(result.pt.high downto 10) := lo.ppn(lo.ppn.high downto 10)
    result.mat := lo.mat
    result.ex := ex_list.map(_._1).reduce(_||_)
    result.esubc := (if(isFetch) U(0) else Mux(is_ade,LISA.Ex.ADEM.subc,U(0)))
    result.ecode := U(0)
    var blk:WhenContext=null
    val last = ex_list.length - 1
    for(i <- 0 to last){
        val (is,ex) = ex_list(i)
        if(i == 0)blk = when(is){result.ecode := ex.code}
        else blk = blk.elsewhen(is){result.ecode := ex.code}
    }
    blk.otherwise{result.ecode := U(0)}
}
class MmuLookup(val r:LookupReq
            ,val cam:Vec[TlbEntry]
            ,val csr:TlbCsr) extends Area{
    val dam = DaMLookup(r,csr.crmd)
    val dmw = DmwLookup(r,csr )
    val tlb = TlbLookup(r,csr,cam)

    def t_found = Bool()
    def t_mat   = UInt(2 bits)
    def t_pt    = UInt(LISA.w_ppn bits)

    val found = if(r.valid == null) t_found else Reg(t_found)
    val mat   = if(r.valid == null) t_mat   else Reg(t_mat  )
    val pt    = if(r.valid == null) t_pt    else Reg(t_pt   )
    def select:Unit = {
        when(dam.found){
            found := True
            mat   := dam.result.mat
            pt    := dam.result.pt
        }.elsewhen(dmw.found){
            found := True
            mat   := dmw.result.mat
            pt    := dmw.result.pt
        }.otherwise{
            found := False
            mat   := U(0)
            pt    := U(0)
        }
    }
    if(r.valid != null){
        when(r.valid){
            select
        }
    }
    else{
        select
    }
    val result = new TransResult
    when(found){
        result := new TransResult assignedFrom(pt,mat)
    }.otherwise{
        result := tlb.result
    }
}
case class LookupReq(srch:TlbSrch,csr:TlbCsr) extends Bundle{
    val isFetch:Boolean = srch.isFetch
    val noValid:Boolean = srch.valid == null
    val va   :UInt = if(isFetch) srch.va else LISA.GPR
    val asid :UInt = if(isFetch) csr.asid.asid else UInt(LISA.w_asid bits)
    val st   :Bool = if(isFetch) null else srch.st
    val valid:Bool = if(noValid) null else Bool()
    if(!isFetch){
        va    := srch.va
        asid  := csr.asid.asid
    }
    if(!noValid)valid := srch.valid
}
class Tlb extends Component{
    val io = new Bundle{
        val csr:TlbCsr   = in(new TlbCsr)
        val isrch:ITlbSrch = slave(new ITlbSrch)
        val dsrch:DTlbSrch = slave(new DTlbSrch)
        val rw  :TlbRW = in(new TlbRW)
        val inv :InvTlb = in(new InvTlb)
        val srch:TlbSrchRes = out(new TlbSrchRes)
        val rd  :TlbRdRes = out(new TlbRdRes)
    }
    val w:Int = io.csr.tlbidx.index.getWidth
    val n:Int = 1 << w
    val cam = Reg(Vec(new TlbEntry,n))

    val itlbreq = LookupReq(io.isrch,io.csr)
    val dtlbreq = LookupReq(io.dsrch,io.csr)
    when(io.rw.srch){
        dtlbreq.valid := True
        dtlbreq.va := io.csr.tlbehi.vppn << 13
    }
    when(io.inv.may){
        dtlbreq.va := io.inv.va
        dtlbreq.asid := io.inv.asid 
    }
    val ilookup = new MmuLookup(itlbreq,cam,io.csr)
    val dlookup = new MmuLookup(dtlbreq,cam,io.csr)

    io.isrch.res := ilookup.result
    io.dsrch.res := dlookup.result 

    io.srch.valid := RegNext(io.rw.srch)
    io.rd  .valid := io.rw.rd
    io.srch.assignedFrom(dlookup)
    io.rd  .readFrom(cam,io.csr)

    val wd_lo:Vec[CsrTLBELONoG] = Vec(io.csr.tlbelo.map((lo:CsrTLBELONoG) => lo))
    val wd_hi:TlbEHi = io.csr.toTlbEHi
    val lfsr = RegInit(U(1,14 bits))
    when(io.rw.fill){
        lfsr := Mux(lfsr.msb,U(0x21,14 bits),U(0)) ^ (lfsr |<< 1)
    }
    val wa:UInt = Mux(io.rw.fill,lfsr(0,w bits),io.csr.tlbidx.index)
    val inv_mask:Bits = io.inv.getInvMask(dlookup.tlb)
    for(i <- 0 until cam.length){
        when(io.rw.en && wa === U(i)){
            cam(i).hi := wd_hi
            cam(i).lo := wd_lo
        }.elsewhen(io.inv.en && inv_mask(i)){
            cam(i).hi.e := False
        }
    }
}
